"""
MeanShift算法实现模块

该模块实现了基于颜色直方图的MeanShift目标跟踪算法。
主要功能：
1. 目标颜色模型建立
2. 候选区域搜索
3. 相似度计算
4. 迭代优化目标位置

作者：目标跟踪系统
日期：2024
"""

import cv2
import numpy as np
from typing import Tuple, Optional


class MeanShiftTracker:
    """
    MeanShift目标跟踪器类
    
    基于颜色直方图的非参数密度估计方法实现目标跟踪。
    使用HSV颜色空间建立目标模型，通过Bhattacharyya系数计算相似度。
    
    参数：
        max_iter: 最大迭代次数，默认为20
        tolerance: 收敛阈值，默认为1.0像素
        hist_bins: 直方图分箱数，默认为16
    """
    
    def __init__(self, max_iter: int = 20, tolerance: float = 1.0, hist_bins: int = 16):
        self.max_iter = max_iter
        self.tolerance = tolerance
        self.hist_bins = hist_bins
        self.target_hist = None
        self.target_center = None
        self.target_size = None
        
    def initialize_target(self, image: np.ndarray, bbox: Tuple[int, int, int, int]) -> None:
        """
        初始化目标模型
        
        从给定的边界框中提取目标区域，建立颜色直方图模型。
        
        参数：
            image: 输入图像，BGR格式
            bbox: 目标边界框 (x, y, width, height)
        """
        x, y, w, h = bbox
        self.target_center = (x + w // 2, y + h // 2)
        self.target_size = (w, h)
        
        # 提取目标区域
        roi = image[y:y+h, x:x+w]
        
        # 转换到HSV颜色空间
        hsv_roi = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
        
        # 创建椭圆掩码，减少背景干扰
        mask = self._create_ellipse_mask(w, h)
        
        # 计算目标直方图
        self.target_hist = cv2.calcHist([hsv_roi], [0, 1], mask, 
                                       [self.hist_bins, self.hist_bins], 
                                       [0, 180, 0, 256])
        
        # 归一化直方图
        cv2.normalize(self.target_hist, self.target_hist, 0, 255, cv2.NORM_MINMAX)
        
    def _create_ellipse_mask(self, width: int, height: int) -> np.ndarray:
        """
        创建椭圆形掩码
        
        用于减少边界区域的背景干扰，突出目标中心区域。
        
        参数：
            width: 掩码宽度
            height: 掩码高度
            
        返回：
            椭圆形掩码，uint8类型
        """
        mask = np.zeros((height, width), dtype=np.uint8)
        center = (width // 2, height // 2)
        axes = (width // 2, height // 2)
        cv2.ellipse(mask, center, axes, 0, 0, 360, 255, -1)
        return mask
        
    def track(self, image: np.ndarray, search_window: Optional[Tuple[int, int, int, int]] = None) -> Tuple[int, int, int, int]:
        """
        在新帧中跟踪目标
        
        使用MeanShift算法在搜索窗口中寻找与目标模型最相似的区域。
        
        参数：
            image: 当前帧图像，BGR格式
            search_window: 搜索窗口 (x, y, width, height)，如果为None则使用默认窗口
            
        返回：
            更新后的目标位置 (x, y, width, height)
        """
        if self.target_hist is None:
            raise ValueError("目标模型未初始化，请先调用initialize_target方法")
            
        # 设置搜索窗口
        if search_window is None:
            x = max(0, self.target_center[0] - self.target_size[0])
            y = max(0, self.target_center[1] - self.target_size[1])
            w = min(image.shape[1] - x, self.target_size[0] * 2)
            h = min(image.shape[0] - y, self.target_size[1] * 2)
            search_window = (x, y, w, h)
            
        # 转换到HSV颜色空间
        hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        
        # 执行MeanShift迭代
        new_center = self._meanshift_iteration(hsv_image, search_window)
        
        # 更新目标中心
        self.target_center = new_center
        
        # 返回更新后的边界框
        x = new_center[0] - self.target_size[0] // 2
        y = new_center[1] - self.target_size[1] // 2
        return (x, y, self.target_size[0], self.target_size[1])
        
    def _meanshift_iteration(self, hsv_image: np.ndarray, search_window: Tuple[int, int, int, int]) -> Tuple[int, int]:
        """
        执行MeanShift迭代过程
        
        在搜索窗口中迭代寻找最佳匹配位置。
        
        参数：
            hsv_image: HSV格式的图像
            search_window: 搜索窗口 (x, y, width, height)
            
        返回：
            新的目标中心位置 (x, y)
        """
        sx, sy, sw, sh = search_window
        current_center = self.target_center
        
        for _ in range(self.max_iter):
            # 计算当前位置的候选区域
            cx, cy = current_center
            roi_x = max(sx, cx - self.target_size[0] // 2)
            roi_y = max(sy, cy - self.target_size[1] // 2)
            roi_w = min(sw - (roi_x - sx), self.target_size[0])
            roi_h = min(sh - (roi_y - sy), self.target_size[1])
            
            if roi_w <= 0 or roi_h <= 0:
                break
                
            # 提取候选区域
            roi = hsv_image[roi_y:roi_y+roi_h, roi_x:roi_x+roi_w]
            
            # 计算反向投影
            backproj = cv2.calcBackProject([roi], [0, 1], self.target_hist, 
                                         [0, 180, 0, 256], 1)
            
            # 计算新的中心位置
            moments = cv2.moments(backproj)
            if moments['m00'] != 0:
                new_cx = int(moments['m10'] / moments['m00']) + roi_x
                new_cy = int(moments['m01'] / moments['m00']) + roi_y
            else:
                new_cx, new_cy = current_center
                
            # 检查收敛条件
            distance = np.sqrt((new_cx - cx)**2 + (new_cy - cy)**2)
            if distance < self.tolerance:
                break
                
            current_center = (new_cx, new_cy)
            
        return current_center
        
    def get_similarity(self, image: np.ndarray, bbox: Tuple[int, int, int, int]) -> float:
        """
        计算给定区域与目标模型的相似度
        
        使用Bhattacharyya系数计算相似度，值越大表示越相似。
        
        参数：
            image: 输入图像，BGR格式
            bbox: 候选区域边界框 (x, y, width, height)
            
        返回：
            相似度值，范围[0, 1]
        """
        if self.target_hist is None:
            return 0.0
            
        x, y, w, h = bbox
        if x < 0 or y < 0 or x + w > image.shape[1] or y + h > image.shape[0]:
            return 0.0
            
        # 提取候选区域
        roi = image[y:y+h, x:x+w]
        hsv_roi = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
        
        # 计算候选区域直方图
        candidate_hist = cv2.calcHist([hsv_roi], [0, 1], None, 
                                     [self.hist_bins, self.hist_bins], 
                                     [0, 180, 0, 256])
        cv2.normalize(candidate_hist, candidate_hist, 0, 255, cv2.NORM_MINMAX)
        
        # 计算Bhattacharyya系数
        similarity = cv2.compareHist(self.target_hist, candidate_hist, cv2.HISTCMP_BHATTACHARYYA)
        
        # 转换为相似度（值越大越相似）
        return 1.0 - similarity 